#include <iostream>
#include <vector>
#include <string>
#include <fstream>
#include "grid_io.hpp"
#include <chrono>

#define MIN3(a, b, c) (a < b ? (a < c ? a : c) : (b < c ? b : c))
#define data_t double

template <typename T, typename V>
void calProps(int nxa, int nya, int nza, int nxb, int nyb, int nzb, T gx, T gy, T gz, \
    std::vector<V>& youngmodsa, std::vector<V>& poissonrsa, std::vector<V>& youngmodsb, std::vector<V>& poissonrsb)
{
    // Calculate the size of a single grid cell for the coarse grid (grid A).
    T dxa = gx / nxa, dya = gy / nya, dza = gz / nza;
    // Calculate the size of a single grid cell for the fine grid (grid B).
    T dxb = gx / nxb, dyb = gy / nyb, dzb = gz / nzb;

    std::cout << "dxa: " << dxa << " dya: " << dya << " dza: " << dza << std::endl;
    std::cout << "dxb: " << dxb << " dyb: " << dyb << " dzb: " << dzb << std::endl;

    // Calculate the volume of a single grid cell for grid A and grid B.
    T va = dxa * dya * dza, vb = dxb * dyb * dzb;

    // Loop over all grid cells in the fine grid (grid B).
    // #pragma omp parallel for collapse(3)
    for (int k = 0; k < nzb; k++) // Loop over the Z dimension.
    {
        for (int j = 0; j < nyb; j++) // Loop over the Y dimension.
        {
            for (int i = 0; i < nxb; i++) // Loop over the X dimension.
            {
                // Calculate the spatial range of the current fine grid cell (grid B).
                T xmin = i * dxb, xmax = (i + 1) * dxb; // X boundaries.
                T ymin = j * dyb, ymax = (j + 1) * dyb; // Y boundaries.
                T zmin = k * dzb, zmax = (k + 1) * dzb; // Z boundaries.

                // Determine the range of indices in the coarse grid (grid A) that overlap with the current fine grid cell.
                int idxamin = static_cast<int>(xmin / dxa), idxamax = std::min(static_cast<int>(xmax / dxa), nxa - 1);
                int idyamin = static_cast<int>(ymin / dya), idyamax = std::min(static_cast<int>(ymax / dya), nya - 1);
                int idzamin = static_cast<int>(zmin / dza), idzamax = std::min(static_cast<int>(zmax / dza), nza - 1);

                // Initialize accumulators for the interpolated properties.
                V youngmodsum = 0, poissonrsum = 0;

                // Loop through all overlapping coarse grid cells (grid A).
                for (int kk = idzamin; kk <= idzamax; kk++) // Z dimension.
                {
                    for (int jj = idyamin; jj <= idyamax; jj++) // Y dimension.
                    {
                        for (int ii = idxamin; ii <= idxamax; ii++) // X dimension.
                        {
                            // Calculate the overlap dimensions between the current fine grid cell and a coarse grid cell.
                            T lx = MIN3((ii + 1) * dxa - xmin, xmax - ii * dxa, dxa); // Overlap in X.
                            T ly = MIN3((jj + 1) * dya - ymin, ymax - jj * dya, dya); // Overlap in Y.
                            T lz = MIN3((kk + 1) * dza - zmin, zmax - kk * dza, dza); // Overlap in Z.

                            // Calculate the overlap volume.
                            T v = lx * ly * lz;

                            // Accumulate the weighted contributions of the properties from the overlapping coarse grid cell.
                            youngmodsum += youngmodsa[kk * nya * nxa + jj * nxa + ii] * v / vb;
                            poissonrsum += poissonrsa[kk * nya * nxa + jj * nxa + ii] * v / vb;
                        }
                    }
                }

                // Assign the interpolated properties to the current fine grid cell (grid B).
                youngmodsb[k * nyb * nxb + j * nxb + i] = youngmodsum;
                poissonrsb[k * nyb * nxb + j * nxb + i] = poissonrsum;
            }
        }
    }
}


int main()
{
    int nxa, nya, nza;
    data_t gx, gy, gz;
    std::string filenameA_coord = "../../../../../data/Grid_A.out";
    auto start_time = std::chrono::high_resolution_clock::now();
    serial_read_GridSize<data_t>(filenameA_coord, nxa, nya, nza, gx, gy, gz, true);
    std::cout << "Grid A size: " << nxa << " " << nya << " " << nza << std::endl;
    std::cout << "Grid range: " << gx << " " << gy << " " << gz << std::endl;

    std::string filenameA_porps = "../../../../../data/A_porps.txt";
    std::vector<data_t> youngmodsa(nxa * nya * nza);
    std::vector<data_t> poissonrsa(nxa * nya * nza);
    serial_read_Porps<data_t>(filenameA_porps, youngmodsa, poissonrsa, nxa, nya, nza);

    int nxb, nyb, nzb;
    std::string filenameB_coord = "../../../../../data/Grid_B.out";
    serial_read_GridSize<data_t>(filenameB_coord, nxb, nyb, nzb, gx, gy, gz, false);
    std::cout << "Grid B size: " << nxb << " " << nyb << " " << nzb << std::endl;
    auto end_time = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double> elapsed = end_time - start_time;
    std::chrono::duration<double> total_time = elapsed;
    std::cout << "Time to read: " << elapsed.count() << "s\n";

    std::vector<data_t> youngmodsb(nxb * nyb * nzb);
    std::vector<data_t> poissonrsb(nxb * nyb * nzb);
    start_time = std::chrono::high_resolution_clock::now();
    calProps<data_t, data_t>(nxa, nya, nza, nxb, nyb, nzb, gx, gy, gz, youngmodsa, poissonrsa, youngmodsb, poissonrsb);
    end_time = std::chrono::high_resolution_clock::now();
    elapsed = end_time - start_time;
    total_time += elapsed;
    std::cout << "Time to calculate porps: " << elapsed.count() << "s\n";

    std::string filenameB_porps = "../../../../../data/B_porps.txt";
    start_time = std::chrono::high_resolution_clock::now();
    serial_write_Props<data_t>(filenameB_porps, youngmodsb, poissonrsb, nxb, nyb, nzb);
    end_time = std::chrono::high_resolution_clock::now();
    elapsed = end_time - start_time;
    total_time += elapsed;
    std::cout << "Time to write porps: " << elapsed.count() << "s\n";

    std::cout << "Total time: " << total_time.count() << "s\n";

    return 0;
}